# -*- coding: utf-8 -*-
import os
import sys
BASE_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.join(BASE_DIR, '..'))
import argparse
import torch
import cv2
import time
import albumentations as A
import numpy as np
from models.build_BiSeNet import BiSeNet


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

parser = argparse.ArgumentParser(description='Training')
parser.add_argument('--path_checkpoint', default=r"G:\img_seg\seg_baseline\checkpoint_best.pkl",
                    help="path to your dataset")
parser.add_argument('--path_img', default=r"E:\seg_dataset\Portrait-dataset-2000\dataset\testing\00006.png",
                    help="path to your dataset")
parser.add_argument('--data_root_dir', default=r"E:\seg_dataset\Portrait-dataset-2000\dataset",
                    help="path to your dataset")
args = parser.parse_args()

if __name__ == '__main__':

    # path_img = args.path_img
    path_img = r"E:\seg_dataset\Portrait-dataset-2000\dataset\testing\00060.png"
    in_size = 512  # 224， 448 ， 336 ， 1024
    norm_mean = (0.5, 0.5, 0.5)  # 比imagenet的mean效果好
    norm_std = (0.5, 0.5, 0.5)
    # step1: 数据预处理
    transform = A.Compose([
        A.Resize(width=in_size, height=in_size),
        A.Normalize(norm_mean, norm_std),
    ])
    img_bgr = cv2.imread(path_img)     # 读取图片
    img_rgb = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2RGB)  # opencv默认bgr，需要转换为rgb
    # cv2.cvtColor() 用于在不同颜色空间之间转换图像。
    # albumentations库中读取为RGB
    transformed = transform(image=img_rgb, mask=img_rgb)  # 预处理
    img_rgb = transformed['image']  # 取出图片

    img_tensor = torch.tensor(np.array(img_rgb), dtype=torch.float).permute(2, 0, 1) # 转张量，hwc--> chw
    # 通常会先使用 np.array() 函数对其进行复制，以避免意外修改原始数据。
    # permute(2, 0, 1)：改变张量的顺序
    img_tensor.unsqueeze_(0)  # chw --> bchw
    img_tensor = img_tensor.to(device)

    # step2: 模型加载
    """
    在 PyTorch 中，模型的权重是以张量的形式存储在模型实例的参数中的。
    load_state_dict 方法用于将预训练的权重加载到模型实例中。
    但是，在调用 load_state_dict 方法之前，必须先创建一个模型实例，以便有地方存储加载的权重。
    """
    model = BiSeNet(num_classes=1, context_path="resnet101")
    checkpoint = torch.load(args.path_checkpoint)
    model.load_state_dict(checkpoint["model_state_dict"])
    model.to(device)
    model.eval()

    # step3: 推理
    with torch.no_grad():
        # 因cpu/gpu是异步的，因此用cpu对gpu计时，应当等待gpu准备好再即使该操作是等待GPU全部执行结束，CPU才可以读取时间信息。
        torch.cuda.synchronize() # 等待gpu处理完毕
        s = time.time()
        outputs = model(img_tensor)
        torch.cuda.synchronize()  # 该操作是等待GPU全部执行结束，CPU才可以读取时间信息。
        print("{:.4f}s".format(time.time() - s))
        outputs = torch.sigmoid(outputs).squeeze(1)
        pre_label = outputs.data.cpu().numpy()[0]  # 取batch的第一个

    # step4：后处理显示图像
    background = np.zeros_like(img_bgr, dtype=np.uint8)

    background[:] = 255
    alpha_bgr = pre_label
    """
    窗口显示RGB格式的图片
    1. 使用 cv2.cvtColor() 函数将其转换为 BGR 格式
    2. 使用 cv2.imshow() 显示，BGR格式输出
    """

    alpha_bgr = cv2.cvtColor(alpha_bgr, cv2.COLOR_GRAY2BGR)
    h, w, c = img_bgr.shape # 获取 img_bgr 的形状，分别赋值给 h、w 和 c
    alpha_bgr = cv2.resize(alpha_bgr, (w, h)) # 将 alpha_bgr 调整为与 img_bgr 相同的大小，宽在前高在后
    # alpha 混合是一种常用的图像混合技术，用于将两幅图像混合在一起。
    # 它的基本思想是对每个像素，根据其 alpha 值来确定两幅图像中对应像素的权重。

    # fusion
    # 背景与人像混合，使用 alpha_bgr 作为 alpha 通道。
    result = np.uint8(img_bgr * alpha_bgr + background * (1 - alpha_bgr))
    # 将输入图像和结果图像水平连接
    out_img = np.concatenate([img_bgr, result], axis=1)

    cv2.imshow("result", out_img)
    cv2.waitKey() # 等待用户按键，这样 imshow 显示的图像才不会立刻消失



